import pymysql
import logging
from typing import Dict, List

class MySQLBatchModifier:
    def __init__(self, host: str, user: str, password: str, database: str,
                 charset: str = 'utf8mb4', dry_run: bool = False):
        self.config = {
            'host': host,
            'user': user,
            'password': password,
            'database': database,
            'charset': charset,
            'cursorclass': pymysql.cursors.DictCursor
        }
        self.dry_run = dry_run
        logging.basicConfig(level=logging.INFO)

    def _get_connection(self):
        return pymysql.connect(**self.config)

    def get_all_tables(self) -> List[str]:
        """获取当前数据库所有表名"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute(
                    "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES "
                    "WHERE TABLE_SCHEMA = DATABASE()")
                return [row['TABLE_NAME'] for row in cursor.fetchall()]
        finally:
            conn.close()

    def column_exists(self, table: str, column: str) -> bool:
        """检查字段是否存在"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute(
                    "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS "
                    "WHERE TABLE_SCHEMA = DATABASE() "
                    "AND TABLE_NAME = %s AND COLUMN_NAME = %s",
                    (table, column))
                return cursor.fetchone() is not None
        finally:
            conn.close()

    def generate_alter_sql(self, table: str, modification: Dict) -> str:
        """生成单个ALTER语句"""
        alter_sql = f"ALTER TABLE `{table}` "
        
        # 字段重命名逻辑
        if 'rename' in modification:
            old_name = modification['rename']['old_name']
            new_name = modification['rename']['new_name']
            if not self.column_exists(table, old_name):
                return None
            alter_sql += f"CHANGE COLUMN `{old_name}` `{new_name}`"
            
        # 字段修改逻辑
        else:
            column = modification['column']
            if not self.column_exists(table, column):
                return None
            alter_sql += f"MODIFY COLUMN `{column}`"

        # 拼接类型定义
        alter_sql += f" {modification['type']}"

        # 处理空约束
        if 'nullable' in modification:
            alter_sql += " NULL" if modification['nullable'] else " NOT NULL"

        # 处理默认值
        if 'default' in modification:
            default = modification['default']
            alter_sql += f" DEFAULT {default}" if default else " DROP DEFAULT"

        return alter_sql

    def execute_global_modify(self, modifications: List[Dict]):
        """全库执行字段修改"""
        conn = None
        try:
            conn = self._get_connection()
            cursor = conn.cursor()
            all_tables = self.get_all_tables()

            if not self.dry_run:
                conn.begin()

            total_modified = 0
            for table in all_tables:
                for mod in modifications:
                    # 生成SQL前自动检查字段存在性
                    sql = self.generate_alter_sql(table, mod)
                    if not sql:
                        continue

                    try:
                        logging.info(f"Processing {table}: {sql}")
                        if not self.dry_run:
                            cursor.execute(sql)
                            total_modified += 1
                    except Exception as e:
                        logging.warning(f"Modify failed on {table}: {str(e)}")
                        if not self.dry_run:
                            conn.rollback()
                            conn.begin()  # 重启事务

            if self.dry_run:
                logging.info("Dry run completed. No changes were made.")
            else:
                conn.commit()
                logging.info(f"Successfully modified {total_modified} columns.")

        except Exception as e:
            logging.error(f"Fatal error: {str(e)}")
            if conn and not self.dry_run:
                conn.rollback()
            raise
        finally:
            if conn:
                conn.close()

if __name__ == "__main__":
    # 配置示例
    modifier = MySQLBatchModifier(
        host='123.56.171.51',
        user='root',
        password='lgmi6395',
        database='jg_base_new',
        dry_run=False  # 生产环境设为False
    )

    # 全库修改规则定义
    modification_rules = [
        # 重名字段规则
        {
            'rename': {
                'old_name': 'CREATE_BY',
                'new_name': 'create_user'
            },
            'type': 'varchar(30)'
        },
        {
            'rename': {
                'old_name': 'CREATE_BY_NAME',
                'new_name': 'create_user_name'
            },
            'type': 'varchar(50)'
        },
        {
            'rename': {
                'old_name': 'UPDATE_BY',
                'new_name': 'change_user'
            },
            'type': 'varchar(30)'
        },
        {
            'rename': {
                'old_name': 'UPDATE_BY_NAME',
                'new_name': 'change_user_name'
            },
            'type': 'varchar(50)'
        },
        {
            'rename': {
                'old_name': 'CREATE_TIME',
                'new_name': 'create_date'
            },
            'type': 'datetime'
        },
        {
            'rename': {
                'old_name': 'UPDATE_TIME',
                'new_name': 'change_date'
            },
            'type': 'datetime'
        }
    ]

    modifier.execute_global_modify(modification_rules)
